import socket
import threading
import json
import base64
from datetime import datetime
import os


class ChatServer:
    def __init__(self, host='localhost', port=12345):
        self.host = host
        self.port = port
        self.clients = {}  # {socket: {'username': str, 'room': str}}
        self.rooms = {}  # {room_name: [socket1, socket2, ...]}
        self.server_socket = None

        self.logs_dir = 'chat_logs'
        if not os.path.exists(self.logs_dir):
            os.makedirs(self.logs_dir)
            print(f"创建聊天记录目录: {self.logs_dir}")

    def start_server(self):
        """启动服务器"""
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        try:
            self.server_socket.bind((self.host, self.port))
            self.server_socket.listen(5)
            print(f"聊天服务器启动在 {self.host}:{self.port}")

            while True:
                client_socket, address = self.server_socket.accept()
                print(f"新连接来自: {address}")

                client_thread = threading.Thread(
                    target=self.handle_client,
                    args=(client_socket,)
                )
                client_thread.daemon = True
                client_thread.start()

        except Exception as e:
            print(f"服务器错误: {e}")
        finally:
            if self.server_socket:
                self.server_socket.close()

    def receive_full_message(self, client_socket):
        """接收完整消息，支持大文件传输"""
        try:
            # 先接收消息长度
            length_data = b''
            while len(length_data) < 4:
                chunk = client_socket.recv(4 - len(length_data))
                if not chunk:
                    return None
                length_data += chunk

            message_length = int.from_bytes(length_data, byteorder='big')

            # 接收完整消息
            message_data = b''
            while len(message_data) < message_length:
                chunk = client_socket.recv(min(4096, message_length - len(message_data)))
                if not chunk:
                    return None
                message_data += chunk

            return message_data.decode('utf-8')
        except Exception as e:
            print(f"接收消息错误: {e}")
            return None

    def handle_client(self, client_socket):
        """处理客户端连接"""
        try:
            while True:
                data = self.receive_full_message(client_socket)
                if not data:
                    break

                try:
                    message = json.loads(data)
                    self.process_message(client_socket, message)
                except json.JSONDecodeError as e:
                    print(f"收到无效的JSON数据: {e}")

        except Exception as e:
            print(f"处理客户端时出错: {e}")
        finally:
            self.disconnect_client(client_socket)

    def process_message(self, client_socket, message):
        """处理不同类型的消息"""
        msg_type = message.get('type')

        if msg_type == 'join':
            self.handle_join(client_socket, message)
        elif msg_type == 'chat':
            self.handle_chat_message(client_socket, message)
        elif msg_type == 'image':
            self.handle_image_message(client_socket, message)
        elif msg_type == 'emoji':
            self.handle_emoji_message(client_socket, message)
        elif msg_type == 'leave':
            self.handle_leave(client_socket, message)
        elif msg_type == 'refresh_users':
            self.handle_refresh_users(client_socket, message)
        elif msg_type == 'private_chat':
            self.handle_private_chat_message(client_socket, message)
        elif msg_type == 'private_image':
            self.handle_private_image_message(client_socket, message)
        elif msg_type == 'private_emoji':
            self.handle_private_emoji_message(client_socket, message)

    def handle_join(self, client_socket, message):
        """处理用户加入房间"""
        username = message.get('username')
        room = message.get('room', 'general')

        # 保存客户端信息
        self.clients[client_socket] = {
            'username': username,
            'room': room
        }

        # 添加到房间
        if room not in self.rooms:
            self.rooms[room] = []
        self.rooms[room].append(client_socket)

        # 通知房间内其他用户
        timestamp = datetime.now().strftime('%H:%M:%S')
        join_message = {
            'type': 'system',
            'message': f"{username} 加入了房间",
            'timestamp': timestamp,
            'room': room
        }

        self.save_message_to_log(room, 'system', 'SYSTEM', f"{username} 加入了房间", timestamp)

        self.broadcast_to_room(room, join_message, exclude=client_socket)

        # 向房间内所有用户广播用户列表，而不是只发给新用户
        users_in_room = [
            self.clients[sock]['username']
            for sock in self.rooms[room]
            if sock in self.clients
        ]
        user_list_message = {
            'type': 'user_list',
            'users': users_in_room,
            'room': room
        }
        self.broadcast_to_room(room, user_list_message)

        print(f"{username} 加入房间 {room}")

    def handle_chat_message(self, client_socket, message):
        """处理文字消息"""
        if client_socket not in self.clients:
            return

        client_info = self.clients[client_socket]
        timestamp = datetime.now().strftime('%H:%M:%S')
        chat_message = {
            'type': 'chat',
            'username': client_info['username'],
            'message': message.get('message'),
            'timestamp': timestamp,
            'room': client_info['room']
        }

        self.save_message_to_log(
            client_info['room'],
            'chat',
            client_info['username'],
            message.get('message'),
            timestamp
        )

        self.broadcast_to_room(client_info['room'], chat_message)

    def handle_image_message(self, client_socket, message):
        """处理图片消息"""
        if client_socket not in self.clients:
            return

        client_info = self.clients[client_socket]

        try:
            image_data = message.get('image_data')
            if not image_data:
                print("收到空的图片数据")
                return

            # 验证base64数据
            try:
                base64.b64decode(image_data)
            except Exception as e:
                print(f"无效的base64图片数据: {e}")
                return

            timestamp = datetime.now().strftime('%H:%M:%S')
            filename = message.get('filename', 'image.jpg')

            image_message = {
                'type': 'image',
                'username': client_info['username'],
                'image_data': image_data,
                'filename': filename,
                'timestamp': timestamp,
                'room': client_info['room']
            }

            self.save_message_to_log(
                client_info['room'],
                'image',
                client_info['username'],
                filename,
                timestamp
            )

            print(f"{client_info['username']} 发送图片: {filename}")
            self.broadcast_to_room(client_info['room'], image_message)

        except Exception as e:
            print(f"处理图片消息错误: {e}")

    def handle_emoji_message(self, client_socket, message):
        """处理表情包消息"""
        if client_socket not in self.clients:
            return

        client_info = self.clients[client_socket]
        timestamp = datetime.now().strftime('%H:%M:%S')
        emoji = message.get('emoji')

        emoji_message = {
            'type': 'emoji',
            'username': client_info['username'],
            'emoji': emoji,
            'timestamp': timestamp,
            'room': client_info['room']
        }

        self.save_message_to_log(
            client_info['room'],
            'emoji',
            client_info['username'],
            emoji,
            timestamp
        )

        self.broadcast_to_room(client_info['room'], emoji_message)

    def handle_leave(self, client_socket, message):
        """处理用户离开"""
        self.disconnect_client(client_socket)

    def broadcast_to_room(self, room, message, exclude=None):
        """向房间内所有用户广播消息"""
        if room not in self.rooms:
            return

        disconnected_clients = []

        for client_socket in self.rooms[room]:
            if exclude and client_socket == exclude:
                continue

            try:
                self.send_to_client(client_socket, message)
            except:
                disconnected_clients.append(client_socket)

        # 清理断开的连接
        for client_socket in disconnected_clients:
            self.disconnect_client(client_socket)

    def send_to_client(self, client_socket, message):
        """向特定客户端发送消息"""
        try:
            data = json.dumps(message, ensure_ascii=False)
            data_bytes = data.encode('utf-8')

            length_bytes = len(data_bytes).to_bytes(4, byteorder='big')
            client_socket.send(length_bytes + data_bytes)

        except Exception as e:
            print(f"发送消息失败: {e}")
            raise

    def disconnect_client(self, client_socket):
        """断开客户端连接"""
        if client_socket in self.clients:
            client_info = self.clients[client_socket]
            username = client_info['username']
            room = client_info['room']

            # 从房间移除
            if room in self.rooms and client_socket in self.rooms[room]:
                self.rooms[room].remove(client_socket)

                # 通知其他用户
                timestamp = datetime.now().strftime('%H:%M:%S')
                leave_message = {
                    'type': 'system',
                    'message': f"{username} 离开了房间",
                    'timestamp': timestamp,
                    'room': room
                }

                self.save_message_to_log(room, 'system', 'SYSTEM', f"{username} 离开了房间", timestamp)

                self.broadcast_to_room(room, leave_message)

                users_in_room = [
                    self.clients[sock]['username']
                    for sock in self.rooms.get(room, [])
                    if sock in self.clients
                ]
                user_list_message = {
                    'type': 'user_list',
                    'users': users_in_room,
                    'room': room
                }
                self.broadcast_to_room(room, user_list_message)

            # 清理客户端信息
            del self.clients[client_socket]
            print(f"{username} 断开连接")

        try:
            client_socket.close()
        except:
            pass

    def save_message_to_log(self, room, message_type, username, content, timestamp):
        """保存消息到日志文件"""
        try:
            log_filename = f"{room}_{datetime.now().strftime('%Y-%m-%d')}.txt"
            log_path = os.path.join(self.logs_dir, log_filename)

            # 格式化消息内容
            if message_type == 'chat':
                log_entry = f"[{timestamp}] {username}: {content}\n"
            elif message_type == 'image':
                log_entry = f"[{timestamp}] {username} 发送了图片: {content}\n"
            elif message_type == 'emoji':
                log_entry = f"[{timestamp}] {username}: {content}\n"
            elif message_type == 'system':
                log_entry = f"[{timestamp}] 系统消息: {content}\n"
            else:
                log_entry = f"[{timestamp}] {username}: {content}\n"

            # 写入日志文件
            with open(log_path, 'a', encoding='utf-8') as f:
                f.write(log_entry)

        except Exception as e:
            print(f"保存聊天记录失败: {e}")

    def handle_refresh_users(self, client_socket, message):
        """处理刷新用户列表请求"""
        if client_socket not in self.clients:
            return

        client_info = self.clients[client_socket]
        room = client_info['room']

        # 获取房间内所有用户
        users_in_room = [
            self.clients[sock]['username']
            for sock in self.rooms.get(room, [])
            if sock in self.clients
        ]

        user_list_message = {
            'type': 'user_list',
            'users': users_in_room,
            'room': room
        }

        self.send_to_client(client_socket, user_list_message)

    def handle_private_chat_message(self, client_socket, message):
        """处理私聊文字消息"""
        if client_socket not in self.clients:
            return

        client_info = self.clients[client_socket]
        target_username = message.get('target_user')

        # 查找目标用户
        target_socket = self.find_user_socket(target_username)
        if not target_socket:
            # 发送错误消息给发送者
            error_message = {
                'type': 'system',
                'message': f"用户 {target_username} 不在线",
                'timestamp': datetime.now().strftime('%H:%M:%S'),
                'room': client_info['room']
            }
            self.send_to_client(client_socket, error_message)
            return

        timestamp = datetime.now().strftime('%H:%M:%S')
        private_message = {
            'type': 'private_chat',
            'username': client_info['username'],
            'message': message.get('message'),
            'timestamp': timestamp,
            'target_user': target_username
        }

        # 发送给目标用户
        self.send_to_client(target_socket, private_message)

        # 发送确认给发送者
        confirm_message = {
            'type': 'private_chat_sent',
            'username': client_info['username'],
            'message': message.get('message'),
            'timestamp': timestamp,
            'target_user': target_username
        }
        self.send_to_client(client_socket, confirm_message)

        # 保存私聊记录
        self.save_private_message_to_log(
            client_info['username'],
            target_username,
            'chat',
            message.get('message'),
            timestamp
        )

    def handle_private_image_message(self, client_socket, message):
        """处理私聊图片消息"""
        if client_socket not in self.clients:
            return

        client_info = self.clients[client_socket]
        target_username = message.get('target_user')

        target_socket = self.find_user_socket(target_username)
        if not target_socket:
            error_message = {
                'type': 'system',
                'message': f"用户 {target_username} 不在线",
                'timestamp': datetime.now().strftime('%H:%M:%S'),
                'room': client_info['room']
            }
            self.send_to_client(client_socket, error_message)
            return

        try:
            image_data = message.get('image_data')
            if not image_data:
                return

            timestamp = datetime.now().strftime('%H:%M:%S')
            filename = message.get('filename', 'image.jpg')

            private_image_message = {
                'type': 'private_image',
                'username': client_info['username'],
                'image_data': image_data,
                'filename': filename,
                'timestamp': timestamp,
                'target_user': target_username
            }

            # 发送给目标用户
            self.send_to_client(target_socket, private_image_message)

            # 发送确认给发送者
            confirm_message = {
                'type': 'private_image_sent',
                'username': client_info['username'],
                'filename': filename,
                'timestamp': timestamp,
                'target_user': target_username
            }
            self.send_to_client(client_socket, confirm_message)

            # 保存私聊记录
            self.save_private_message_to_log(
                client_info['username'],
                target_username,
                'image',
                filename,
                timestamp
            )

        except Exception as e:
            print(f"处理私聊图片消息错误: {e}")

    def handle_private_emoji_message(self, client_socket, message):
        """处理私聊表情包消息"""
        if client_socket not in self.clients:
            return

        client_info = self.clients[client_socket]
        target_username = message.get('target_user')

        target_socket = self.find_user_socket(target_username)
        if not target_socket:
            error_message = {
                'type': 'system',
                'message': f"用户 {target_username} 不在线",
                'timestamp': datetime.now().strftime('%H:%M:%S'),
                'room': client_info['room']
            }
            self.send_to_client(client_socket, error_message)
            return

        timestamp = datetime.now().strftime('%H:%M:%S')
        emoji = message.get('emoji')

        private_emoji_message = {
            'type': 'private_emoji',
            'username': client_info['username'],
            'emoji': emoji,
            'timestamp': timestamp,
            'target_user': target_username
        }

        # 发送给目标用户
        self.send_to_client(target_socket, private_emoji_message)

        # 发送确认给发送者
        confirm_message = {
            'type': 'private_emoji_sent',
            'username': client_info['username'],
            'emoji': emoji,
            'timestamp': timestamp,
            'target_user': target_username
        }
        self.send_to_client(client_socket, confirm_message)

        # 保存私聊记录
        self.save_private_message_to_log(
            client_info['username'],
            target_username,
            'emoji',
            emoji,
            timestamp
        )

    def find_user_socket(self, username):
        """根据用户名查找socket连接"""
        for socket, client_info in self.clients.items():
            if client_info['username'] == username:
                return socket
        return None

    def save_private_message_to_log(self, sender, receiver, message_type, content, timestamp):
        """保存私聊消息到日志文件"""
        try:
            # 创建私聊日志目录
            private_logs_dir = os.path.join(self.logs_dir, 'private')
            if not os.path.exists(private_logs_dir):
                os.makedirs(private_logs_dir)

            # 私聊日志文件名：按字母顺序排列用户名
            users = sorted([sender, receiver])
            log_filename = f"{users[0]}_{users[1]}_{datetime.now().strftime('%Y-%m-%d')}.txt"
            log_path = os.path.join(private_logs_dir, log_filename)

            # 格式化消息内容
            if message_type == 'chat':
                log_entry = f"[{timestamp}] {sender} -> {receiver}: {content}\n"
            elif message_type == 'image':
                log_entry = f"[{timestamp}] {sender} -> {receiver} 发送了图片: {content}\n"
            elif message_type == 'emoji':
                log_entry = f"[{timestamp}] {sender} -> {receiver}: {content}\n"
            else:
                log_entry = f"[{timestamp}] {sender} -> {receiver}: {content}\n"

            # 写入日志文件
            with open(log_path, 'a', encoding='utf-8') as f:
                f.write(log_entry)

        except Exception as e:
            print(f"保存私聊记录失败: {e}")


if __name__ == "__main__":
    server = ChatServer()
    try:
        server.start_server()
    except KeyboardInterrupt:
        print("\n服务器关闭")
